

clear


%% data and prior
load('../../true_DGP/priorWSigma.mat')
data=readtable('../../../Data/data_main_sample.csv');
data=data(data.year>=1950&data.year<=2000&isfinite(data.D),:);
ccodes=unique(data.ccode);
years=unique(data.year(data.year>1950));
N=length(ccodes);
T=length(years);
K=0;
y=zeros(N,T); D=y; M=y; X=repmat(y,1,1,K); Y=-ones(N,T);
for n=1:N
    dn=data(data.ccode==ccodes(n),:);
    for t=1:T
        if isempty(setdiff(years(t),dn.year))
            if isfinite(dn.y(years(t)==dn.year))
            	M(n,t)=1; % not missing
                y(n,t)=dn.y(years(t)==dn.year);
                D(n,t)=dn.D(years(t)==dn.year);
                if isfinite(dn.Y(years(t)-1==dn.year))
                    Y(n,t)=dn.Y(years(t)-1==dn.year);
                end
            end
        end
    end
    if Y(n,1)<0
        ybar=max(0,mean(y(n,M(n,:)==1)));
        tmin=find(Y(n,:)>=0,1);
        for t=tmin-1:-1:1
            Y(n,t)=Y(n,t+1)/(1+ybar); % impute 1950 GDPpc using empirical distribution
        end
    end
end
Y=Y(:,1);
clear dn t ybar tmin

% distance measures
 % geographic
dd=readtable('../../../Data/distance_capitals.csv');
Z(:,:,1)=10^(-6)*table2array(dd(:,2:end));
 % genetic
gdist=readtable('../../../Data/gdistance.csv');
gdist.Bosnia=gdist.Yugoslavia; % impute Bosnia as Yugoslavia
for n=1:N
    gdist(43,2+n)=gdist(n,45);
end
gdist(43,45)=gdist(44,46);
gdist=table2array(gdist(:,3:end));
Z(:,:,2)=10^(-3)*gdist;
 % economic
Z(:,:,3)=zeros(N);
for n=1:N-1
    for m=n+1:N
        Z(m,n,3)=(Y(n,1)-Y(m,1))^2/var(Y);
    end
end
Z(:,:,3)=Z(:,:,3)+tril(Z(:,:,3))';
KZ=size(Z,3);

ep_sd=sqrt(diag(Sigma));

% Sparse-grid integration
[sgi_nodes,sgi_weights]=nwspgr('KPN',2,8);

cnames=cell(N,1);
for n=1:N
    a=unique(data.idacr(data.ccode==ccodes(n)));
    cnames{n}=a{1};
end

ccodes=table(ccodes,cnames);

clear a cnames data dd n m gdist Y


%% for Knitro

% Jacobian sparsity pattern:
JPattern=[repmat(eye(N),T,1),ones(N*T,2),ones(N*T,2*N),ones(N*T,N),...
    repmat(eye(N),T,1),zeros(N*T,N),ones(N*T,K),ones(N*T,KZ),eye(N*T)];
JPattern=sparse(JPattern);

% Hessian sparsity pattern
% log-posterior
HPattern=zeros(6*N+2+K+KZ+N*T,6*N+2+K+KZ+N*T);
  % alpha,theta,beta,v,f
HPattern(1:(5*N+2),1:(5*N+2))=eye(5*N+2);
  % vs
HPattern(5*N+2+(1:N),:)=[zeros(N,5*N+2),eye(N),zeros(N,K+KZ),repmat(eye(N),1,T)];
  % kappa
HPattern((6*N+2+K+KZ+1):end,:)=[zeros(N*T,5*N+2),repmat(eye(N),T,1),zeros(N*T,K+KZ),eye(N*T)];
% equilibrium constraints
  % alpha
  HPattern(1:N,:)=HPattern(1:N,:)+...
        [repmat(eye(N),1,2),ones(N,2*N+2),eye(N),zeros(N),ones(N,K+KZ),repmat(eye(N),1,T)];
  % theta
  HPattern(N+(1:2),:)=HPattern(N+1,:)+ones(2,6*N+2+K+KZ+N*T);
  % f
  HPattern(4*N+(1:N),:)=HPattern(4*N+(1:N),:)+...
        [repmat(eye(N),1,2),ones(N,2*N+2),eye(N),zeros(N),ones(N,K+KZ),repmat(eye(N),1,T)];
  % beta,v,xi,gamma
for n=2:4
    HPattern((n-1)*N+2+(1:N),:)=HPattern((n-1)*N+2+(1:N),:)+...
        [ones(N,5*N+2),zeros(N),ones(N,K+KZ+N*T)];
end
HPattern(6*N+2+(1:K),:)=HPattern(6*N+2+(1:K),:)+...
        [ones(K,5*N+2),zeros(K,N),ones(K,K+KZ+N*T)];
HPattern(6*N+2+K+(1:KZ),:)=HPattern(6*N+2+K+(1:KZ),:)+...
        [ones(KZ,5*N+2),zeros(KZ,N),ones(KZ,K+KZ+N*T)];
  % kappa
HPattern((6*N+2+K+KZ+1):end,:)=HPattern((6*N+2+K+KZ+1):end,:)+...
    [repmat([repmat(eye(N),1,2),ones(N,2*N+2),eye(N),zeros(N),ones(N,K+KZ)],T,1),eye(N*T)];
HPattern=1*(HPattern>0);

HPattern=sparse(HPattern);

clear n

% options
options=optimset('Display','off','JacobPattern',JPattern,'HessPattern',HPattern);


%% starting values

% "short" replication
load('robustDistM_replication.mat','MAP')
pars0=MAP;

% % "long" replication
% load('robustDistM_replication.mat','pars0')


%% maximum-a-posteriori estimator
[MAP,lp,exit]=knitromatlab(@(x)log_posterior(x,D,M,X,Z,...
    a0,om_a,th0,om_th,b0A,b0D,om_b,s_v,d_v,f0,om_f,s_vs,d_vs),pars0,[],[],...
  [],[],[-Inf(3*N+2,1);zeros(N,1);-Inf(N,1);zeros(N,1);-Inf(K,1);zeros(KZ,1);-Inf(N*T,1)],...
  Inf(length(pars0),1),@(x)eq_con_jac_robustDistM(x,y,D,M,X,Z,sgi_nodes,sgi_weights,ep_sd,Sigma),[],...
  options,'knitro.opt');


%%
clearvars -except MAP exit
save('MAProbustDistM.mat')



